/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2018-2020.

 * Description: stack repeat check character
 * Author: chenjialong
 * Create: 2018-5-7
 */

#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/errno.h>
#include <linux/stacktrace.h>
#include <linux/fs.h>
#include <linux/uaccess.h>
#include <linux/proc_fs.h>
#include <linux/seq_file.h>
#include <linux/rbtree.h>
#include <linux/string.h>
#include <linux/sched.h>
#include <linux/vmalloc.h>
#include <linux/slab.h>
#include <linux/stackrepeatcheck.h>
#include <asm/traps.h>

#define MAX_CHAR_NAME		256
#define CHAR_ENABLE             '1'
#define CHAR_DISABLE            '0'
#define MAX_LATEST_LENGTH	3
#define MAX_INT			0x7ffffff
#define ENABLE_PROC_LEN		2

struct addr_info {
	unsigned long addr_val[CONFIG_RTOS_MAX_ADDR_INDEX];
	struct task_struct *tsk;
	struct rb_node node;
};

struct addr_managerpool {
	struct rb_root addr_tree;
	unsigned int pool_index;
	struct addr_info *latest_array[MAX_LATEST_LENGTH];
	unsigned int latest_index;
	struct addr_info *addr_info_pool;
	int region_num;
};

static struct addr_managerpool ksmain_addr_managerpool;
static struct addr_managerpool sysrq_addr_managerpool;
static bool check_repeat_stack_enable;

static struct addr_managerpool *search_addr_managerpool(const int region_num)
{
	if (region_num == REGION_FIQ)
		return &ksmain_addr_managerpool;
	else if (region_num == REGION_SYSRQ)
		return &sysrq_addr_managerpool;
	return NULL;
}

static int addr_cmp(const unsigned long *addr_val1,
	const unsigned long *addr_val2)
{
	int i;
	for (i = 0; i < CONFIG_RTOS_MAX_ADDR_INDEX; i++) {
		if (addr_val1[i] < addr_val2[i])
			return -1;
		else if (addr_val1[i] > addr_val2[i])
			return 1;
	}
	return 0;
}

static bool addr_insert_tree(struct addr_managerpool *addr_pool,
	const unsigned long *addr_val, struct task_struct *tsk)
{
	struct rb_node **tmp = &(addr_pool->addr_tree.rb_node);
	struct rb_node *parent = NULL;
	struct addr_info *new_addr_info = NULL;
	struct addr_info *this = NULL;
	int i;
	int ret = 0;

	/* Figure out where to put new node */
	while (*tmp) {
		this = container_of(*tmp, struct addr_info, node);
		ret = addr_cmp(addr_val, this->addr_val);
		parent = *tmp;
		if (ret < 0) {
			tmp = &((*tmp)->rb_left);
		} else if (ret > 0) {
			tmp = &((*tmp)->rb_right);
		} else {
			addr_pool->latest_array[addr_pool->latest_index] = this;
			addr_pool->latest_index++;
			addr_pool->latest_index %= MAX_LATEST_LENGTH;
			printk("Call trace same as TASK [pid=%d comm=%s]\n",
			       this->tsk->pid, this->tsk->comm);
			return false;
		}
	}

	if (addr_pool->pool_index >= CONFIG_RTOS_MAXPOOL_LENGTH)
		return true;

	new_addr_info = &(addr_pool->addr_info_pool[addr_pool->pool_index]);
	addr_pool->pool_index++;

	for (i = 0; i < CONFIG_RTOS_MAX_ADDR_INDEX; i++)
		new_addr_info->addr_val[i] = addr_val[i];

	new_addr_info->tsk = tsk;
	/* Add new node and rebalance tree. */
	rb_link_node(&new_addr_info->node, parent, tmp);
	rb_insert_color(&new_addr_info->node, &(addr_pool->addr_tree));
	return true;
}

static bool insert_addr_new_node(const unsigned long *addr_val,
	struct task_struct *tsk, struct addr_managerpool *addr_pool)
{
	int i;
	struct addr_info *addr_info = NULL;
	unsigned long *tmp_val = NULL;

	for (i = 0; i < MAX_LATEST_LENGTH; i++) {
		addr_info = addr_pool->latest_array[i];
		if (addr_info == NULL)
			continue;
		tmp_val = addr_info->addr_val;
		if (!addr_cmp(addr_val, tmp_val)) {
			printk("Call trace same as TASK [pid=%d comm=%s]\n",
			       addr_info->tsk->pid,
			       addr_info->tsk->comm);
			return false;
		}
	}

	return addr_insert_tree(addr_pool, addr_val, tsk);
}

void clean_stack_addr(const int region_num)
{
	struct addr_managerpool *addr_pool = NULL;
	if (!check_repeat_stack_enable)
		return;

	addr_pool = search_addr_managerpool(region_num);
	if (addr_pool == NULL || addr_pool->addr_info_pool == NULL)
		return;

	memset(addr_pool->addr_info_pool, 0,
	       CONFIG_RTOS_MAXPOOL_LENGTH * sizeof(struct addr_info));
	addr_pool->latest_index = 0;
	addr_pool->pool_index = 0;
	addr_pool->addr_tree = RB_ROOT;
}
EXPORT_SYMBOL_GPL(clean_stack_addr);

bool check_repeat_stack(struct task_struct *tsk, const int region_num)
{
	bool print_stack = true;
	unsigned long address[CONFIG_RTOS_MAX_ADDR_INDEX + 1] = {0};
	unsigned int addr_size;
	struct addr_managerpool *pool = NULL;
	if (!check_repeat_stack_enable)
		return print_stack;

	pool = search_addr_managerpool(region_num);
	if (pool == NULL || pool->addr_info_pool == NULL)
		return print_stack;
#if defined(CONFIG_ARM)
	addr_size = rtos_save_stack_trace(tsk, address,
					  CONFIG_RTOS_MAX_ADDR_INDEX + 1, 0);
#elif defined(CONFIG_ARM64)
	addr_size = stack_trace_save_tsk(tsk, address,
					 CONFIG_RTOS_MAX_ADDR_INDEX + 1, 0);
#endif
	if (addr_size != 0 && addr_size <= CONFIG_RTOS_MAX_ADDR_INDEX)
		print_stack = insert_addr_new_node(address, tsk, pool);

	return print_stack;
}
EXPORT_SYMBOL_GPL(check_repeat_stack);

static int enable_show(struct seq_file *seq, void *offset)
{
	seq_printf(seq, "%d\n", check_repeat_stack_enable);
	return 0;
}

static ssize_t enable_write(struct file *file, const char __user *buf,
	size_t len, loff_t *ppos)
{
	char c;

	if (!capable(CAP_SYS_ADMIN))
		return -EPERM;
	if (len != ENABLE_PROC_LEN)
		return -EINVAL;
	if (copy_from_user(&c, buf, 1))
		return -EFAULT;
	switch (c) {
	case CHAR_ENABLE:
		check_repeat_stack_enable = true;
		break;
	case CHAR_DISABLE:
		check_repeat_stack_enable = false;
		break;
	default:
		return -EINVAL;
	}
	return len;
}

static int enable_open(struct inode *inode, struct file *file)
{
	return single_open(file, enable_show, NULL);
}

static const struct proc_ops proc_enable = {
	.proc_open           = enable_open,
	.proc_read           = seq_read,
	.proc_write          = enable_write,
	.proc_lseek         = seq_lseek,
	.proc_release        = single_release,
};

static int alloc_managerpool(struct addr_managerpool *pool, const int region)
{
	int size = CONFIG_RTOS_MAXPOOL_LENGTH * sizeof(struct addr_info);
	if (pool == NULL)
		return -1;
	pool->addr_info_pool = (struct addr_info *)vmalloc(size);
	if (pool->addr_info_pool == NULL)
		return -1;
	memset(pool->addr_info_pool, 0, size);
	pool->region_num = region;
	pool->addr_tree = RB_ROOT;
	return 0;
}

static int __init stackrepeat_init(void)
{
	struct proc_dir_entry *p = NULL;
	struct proc_dir_entry *parent = NULL;
#ifdef CONFIG_RTOS_KSNAPSHOT
	if (alloc_managerpool(&ksmain_addr_managerpool, REGION_FIQ))
		goto ks_mainfail;
#endif
#ifdef CONFIG_MAGIC_SYSRQ
	if (alloc_managerpool(&sysrq_addr_managerpool, REGION_SYSRQ))
		goto sysrqfail;
#endif
	parent = proc_mkdir("check_repeat_stack", NULL);
	if (parent == NULL) {
		pr_err("proc_create check_repeat_stack failed.\n");
		goto create_dir_failed;
	}
	p = proc_create("enable", (mode_t)0640, parent, &proc_enable);
	if (p == NULL) {
		pr_err("proc_create check_repeat_stack/enable failed.\n");
		goto create_enable_failed;
	}
	return 0;
create_enable_failed:
	remove_proc_entry("check_repeat_stack", NULL);
create_dir_failed:
#ifdef CONFIG_MAGIC_SYSRQ
	vfree(sysrq_addr_managerpool.addr_info_pool);
sysrqfail:
#endif
#ifdef CONFIG_RTOS_KSNAPSHOT
	vfree(ksmain_addr_managerpool.addr_info_pool);
ks_mainfail:
#endif
	return -ENOMEM;
}
late_initcall(stackrepeat_init);
